[Python] ユニットテストでmotoのメソッドをスタブに差し替えたい

こんにちは。サービス開発室の武田です。

Pythonのboto3用モックライブラリmotoを使用することで、ローカルの閉じた環境でboto3を使用するプログラムのユニットテストを実行できます。

バージョンアップを重ねる中でmotoのカバレッジも増えていますが未実装のAPIもあります。そこで次のようなケースでmotoの挙動を変更したいという場合があります。

• 未実装APIのモックを実装したい
• 既存の実装を上書きしたい
• 同じAPIに異なる挙動をさせたい

どの目的でも同じ実装方法が使用できますが、今回は3つ目を例に紹介します。

環境構築

Poetryを使用して検証環境を作ります。

poetry install
poetry add boto3 moto pytest pytest-mock

バージョンは次のようになっていました。

python = "^3.12"
boto3 = "^1.34.70"
moto = "^5.0.3"
pytest = "^8.1.1"
pytest-mock = "^3.14.0"

テスト対象となるコード

今回テスト対象となるプログラムは次のように単純なものです。

import boto3


def main():
    orgs = boto3.client("organizations")

    try:
        return orgs.describe_organization()["Organization"]["MasterAccountId"]
    except orgs.exceptions.AWSOrganizationsNotInUseException:
        return "000000000000"

AWS Organizationsの組織情報を取得し、管理アカウントのアカウントIDを返します。組織に所属していなければ000000000000を返します。

この、「組織に所属しているかどうか」という状態は、「管理アカウントとなるアカウント」で組織を作成し、さらにその組織にアカウントを追加しないと作り出せません。そのため、describe_organizationの実装をすげかえることでユニットテストを実現します。

テストコード

続いて先ほどのプログラムをテストするコードです。

from moto import mock_aws
import botocore
from pytest_mock import MockFixture
import main

orig_api_call = botocore.client.BaseClient._make_api_call


def mock_api_call(self, operation_name, kwarg):
    if operation_name == "DescribeOrganization":
        return {"Organization": {"MasterAccountId": "888888888888"}}

    return orig_api_call(self, operation_name, kwarg)


@mock_aws
def test_main1():
    result = main.main()
    assert result == "000000000000"


@mock_aws
def test_main2(mocker: MockFixture):
    mocker.patch("botocore.client.BaseClient._make_api_call", new=mock_api_call)
    result = main.main()
    assert result == "888888888888"

ポイントとなるのはmock_api_call関数です。mocker.patch関数のnewパラメーターに指定しています。処理内容としては、呼び出されたAPIがDescribeOrganizationであれば、指定したデータを返し、それ以外であれば通常と同じ処理を呼び出すようになっています。

これによってtest_main1は組織に属していない場合のテストケースとなり、test_main2は組織に属している場合のテストケースになります。

実際に実行してみると、テストは成功します。

$ poetry run pytest
collected 2 items

test_main.py ..

まとめ

今回紹介した、mock_api_call関数を定義する方法はboto3のメソッド全般で使用できる汎用的な方法です。これは応用できる範囲が広いのでぜひ覚えておきましょう。